import torch
from torch import nn
from torch.nn import functional as F
from PIL import Image
import numpy as np

import sys
sys.path.append("../../py")
from utils import save_tensor_as_text, save_tensor_as_c_header, float_tensor2Q88
from mobilenetv2_mini import mobilenetv2
from fuse_bn import fuse_module

SAVED_MODEL = "../../py/mobilenetv2_on_cifar10_mini.pth"

if __name__ == '__main__':
    if not len(sys.argv) == 2:
        print("Usage: python {} imagefile".format(sys.argv[0]))
        exit(1)
    img = Image.open(sys.argv[1]).convert('RGB').resize((32, 32))
    img = (np.array(img) / 255.0) * 2.0 - 1.0
    img = torch.tensor(img.transpose((2,0,1))).unsqueeze(0).float()
    net = mobilenetv2()
    net.load_state_dict(torch.load(SAVED_MODEL))
    net = fuse_module(net)
    
    img = img.reshape(3, 32, 32)
    out = net(img).reshape([10, 1, 1])
    
    save_tensor_as_c_header(float_tensor2Q88(img), "img_in")
        